package com.ohos.connect.devicecloud.common.util;

import cn.hutool.core.util.RandomUtil;
import com.alibaba.fastjson.JSON;

import javax.crypto.Cipher;
import javax.crypto.Mac;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.KeySpec;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;

/**
 * @author 毕伟谦
 * @since 2025-06-24
 */

public class CoapEncryptor {
    private static final String PBKDF2_ALGORITHM = "PBKDF2WithHmacSHA256";
    private static final String AES_ALGORITHM = "AES/CBC/PKCS5Padding";
    private static final String HMAC_ALGORITHM = "HmacSHA256";
    private static final int ITERATIONS = 1;
    private static final int KEY_LENGTH = 256; // 32 bytes
    private static final int AES_KEY_LENGTH = 16; // 128 bits
    private static final int IV_LENGTH = 16; // 128 bits
    private static final int MAC_LENGTH = 32; // 256 bits

    /**
     * 执行PBKDF2-HMAC-SHA256密钥派生
     */
    public static byte[] deriveKey(String password, byte[] salt) throws InvalidKeySpecException, NoSuchAlgorithmException {
        KeySpec spec = new PBEKeySpec(password.toCharArray(), salt, ITERATIONS, KEY_LENGTH);
        SecretKeyFactory factory = SecretKeyFactory.getInstance(PBKDF2_ALGORITHM);
        return factory.generateSecret(spec).getEncoded();
    }

    /**
     * 执行AES-CBC-PKCS5Padding加密
     */
    public static byte[] encryptAES(byte[] data, byte[] key, byte[] iv)
            throws Exception {
        SecretKeySpec secretKey = new SecretKeySpec(key, "AES");
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivSpec);
        return cipher.doFinal(data);
    }

    /**
     * AES 128 CBC解密
     *
     * @param encryptedData 密文字符串，
     * @param password      密码
     * @return 明文
     * @throws Exception 解密过程中可能抛出的异常
     */
    public static String decrypt(byte[] encryptedData, String password, String sn1Hex, String sn2Hex) throws Exception {

        // 1. 解析并拼接salt
        byte[] sn1 = hexStringToByteArray(sn1Hex);
        byte[] sn2 = hexStringToByteArray(sn2Hex);
        byte[] salt = concatBytes(sn1, sn2);

        // 2. 生成主密钥
        byte[] derivedKey = deriveKey(password, salt);

        // 3. 提取AES密钥和IV
        byte[] aesKey = Arrays.copyOfRange(derivedKey, 0, AES_KEY_LENGTH);
        byte[] iv = Arrays.copyOfRange(derivedKey, AES_KEY_LENGTH, AES_KEY_LENGTH + IV_LENGTH);
        // 提取原始数据 (不包括coapHeader和 HMAC)
        int dataLength = encryptedData.length - MAC_LENGTH;
        int payLoadLen = dataLength - 4;

        byte[] payload = new byte[payLoadLen];
        System.arraycopy(encryptedData, 4, payload, 0, payLoadLen);

        // 创建密钥规范
        SecretKeySpec secretKeySpec = new SecretKeySpec(aesKey, "AES");
        IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);

        // 初始化解密器
        Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, secretKeySpec, ivParameterSpec);

        // 执行解密
        byte[] decryptedBytes = cipher.doFinal(payload);

        // 返回明文
        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    /**
     * 计算HMAC-SHA256摘要
     */
    public static byte[] calculateHMAC(byte[] data, byte[] secret)
            throws NoSuchAlgorithmException, InvalidKeyException {
        Mac mac = Mac.getInstance(HMAC_ALGORITHM);
        SecretKeySpec secretKey = new SecretKeySpec(secret, HMAC_ALGORITHM);
        mac.init(secretKey);
        return mac.doFinal(data);
    }

    /**
     * 构建CoAP加密消息（含MAC）
     */
    public static byte[] buildEncryptedCoapMessage(String password, String sn1Hex, String sn2Hex,
                                                   byte[] coapHeader, byte[] payload) throws Exception {
        // 1. 解析并拼接salt
        byte[] sn1 = hexStringToByteArray(sn1Hex);
        byte[] sn2 = hexStringToByteArray(sn2Hex);
        byte[] salt = concatBytes(sn1, sn2);

        // 2. 生成主密钥
        byte[] derivedKey = deriveKey(password, salt);

        // 3. 提取AES密钥和IV
        byte[] aesKey = Arrays.copyOfRange(derivedKey, 0, AES_KEY_LENGTH);
        byte[] iv = Arrays.copyOfRange(derivedKey, AES_KEY_LENGTH, AES_KEY_LENGTH + IV_LENGTH);

        // 4. 加密payload
        byte[] encryptedPayload = encryptAES(payload, aesKey, iv);

        // 5. 构建HMAC secret (derivedKey前16字节作为密码)
        String hmacPassword = bytesToHexString(Arrays.copyOfRange(derivedKey, 0, AES_KEY_LENGTH));
        byte[] hmacSecret = deriveKey(hmacPassword, salt);

        // 6. 构建HMAC输入 (coapHeader + encryptedPayload)
        byte[] hmacInput = concatBytes(coapHeader, encryptedPayload);

        // 7. 计算HMAC
        byte[] hmac = calculateHMAC(hmacInput, hmacSecret);

        // 8. 构建最终消息 (coapHeader + encryptedPayload + hmac)
        return concatBytes(hmacInput, hmac);
    }

    /**
     * 验证接收到的 CoAP 消息的完整性
     */
    public static boolean verifyCoapMessage(byte[] message, byte[] coapHeader,
                                            String password, String sn1Hex, String sn2Hex)
            throws NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException {
        // 检查消息长度是否足够
        if (message.length < coapHeader.length + MAC_LENGTH) {
            return false;
        }

        // 提取原始数据 (不包括 HMAC)
        int dataLength = message.length - MAC_LENGTH;
        byte[] data = Arrays.copyOf(message, dataLength);

        // 提取接收到的 HMAC
        byte[] receivedHmac = Arrays.copyOfRange(message, dataLength, message.length);

        // 生成 salt
        byte[] sn1 = hexStringToByteArray(sn1Hex);
        byte[] sn2 = hexStringToByteArray(sn2Hex);
        byte[] salt = concatBytes(sn1, sn2);

        // 2. 生成主密钥
        byte[] derivedKey = deriveKey(password, salt);

        // 5. 构建HMAC secret (derivedKey前16字节作为密码)
        String hmacPassword = bytesToHexString(Arrays.copyOfRange(derivedKey, 0, AES_KEY_LENGTH));
        byte[] hmacSecret = deriveKey(hmacPassword, salt);

        // 6. 计算HMAC
        byte[] hmac = calculateHMAC(data, hmacSecret);

        // 比较两个 HMAC 值
        return Arrays.equals(receivedHmac, hmac);
    }


    // 工具方法: 十六进制字符串转字节数组
    private static byte[] hexStringToByteArray(String s) {
        int len = s.length();
        byte[] data = new byte[len / 2];
        for (int i = 0; i < len; i += 2) {
            data[i / 2] = (byte) ((Character.digit(s.charAt(i), 16) << 4)
                    + Character.digit(s.charAt(i + 1), 16));
        }
        return data;
    }

    // 工具方法: 字节数组转十六进制字符串
    private static String bytesToHexString(byte[] bytes) {
        StringBuilder result = new StringBuilder();
        for (byte b : bytes) {
            result.append(String.format("%02x", b));
        }
        return result.toString();
    }

    // 工具方法: 拼接字节数组
    private static byte[] concatBytes(byte[] a, byte[] b) {
        byte[] result = new byte[a.length + b.length];
        System.arraycopy(a, 0, result, 0, a.length);
        System.arraycopy(b, 0, result, a.length, b.length);
        return result;
    }


    // 示例用法
    public static void main(String[] args) throws InvalidKeySpecException, NoSuchAlgorithmException {

        ArrayList<HashMap<String, Object>> maps = new ArrayList<>();

        HashMap<String, Object> body = new HashMap<>();
        body.put("st", "switch");
        HashMap<String, Object> map = new HashMap<>();
        map.put("on", "1");
        body.put("data", map);
        HashMap<String, Object> body1 = new HashMap<>();
        body1.put("st", "switch");
        HashMap<String, Object> map1 = new HashMap<>();
        map1.put("on", "1");
        body1.put("data", map1);
        HashMap<String, Object> body2 = new HashMap<>();
        body2.put("st", "switch");
        HashMap<String, Object> map2 = new HashMap<>();
        map2.put("on", "1");
        body2.put("data", map2);
        maps.add(body);
        maps.add(body1);
        maps.add(body2);
        //明文
        String plaintext = JSON.toJSONString(maps);


        try {
            // 示例参数
            String password = "yourAuthCodeOrPSK";
            String sn1Hex = "0123456789ABCDEF"; // 8 bytes hex
            String sn2Hex = "FEDCBA9876543210"; // 8 bytes hex
            byte[] coapHeader = new byte[]{0x01, 0x02, 0x03, 0x04}; // 示例CoAP头

            byte[] payload = plaintext.getBytes(StandardCharsets.UTF_8);

            System.out.println("加密前的文件内容:" + plaintext);
            // 构建加密消息
            byte[] encryptedMessage = buildEncryptedCoapMessage(
                    password, sn1Hex, sn2Hex, coapHeader, payload
            );

            System.out.println("加密消息长度: " + encryptedMessage.length + " bytes");
            System.out.println("加密消息Hex: " + bytesToHexString(encryptedMessage));

            boolean isValid = verifyCoapMessage(encryptedMessage, coapHeader, password, sn1Hex, sn2Hex);
            System.out.println("消息完整性验证结果: " + isValid);
            if (isValid) {
                String decrypt = decrypt(encryptedMessage, password, sn1Hex, sn2Hex);
                System.out.println("解密后的文件内容:" + decrypt);
            }


        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static String getRandomSn() {
        // 生成8字节的随机数据
        byte[] randomBytes = RandomUtil.randomBytes(8);

        // 将字节数组转换成十六进制字符串
        // Hutool没有直接提供字节到十六进制字符串的转换方法，所以我们使用Java标准库
        StringBuilder hexString = new StringBuilder(2 * randomBytes.length); // 每个字节变成两个十六进制字符
        for (byte b : randomBytes) {
            // 将每个字节转换成两位的十六进制数并追加到StringBuilder中
            hexString.append(String.format("%02x", b & 0xff));
        }

        return hexString.toString();
    }
}
